#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Minimal Encrypted Chat v3.0 - Stable Version (English)
P-384 ECDH + AES-256-GCM, IPv6 P2P
"""

import socket
import threading
import time
import secrets
import sys
import os
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.kdf.hkdf import HKDF


class SimpleCrypto:
    """Minimal crypto class - P-384 curve, 192-bit security"""
    def __init__(self):
        print("[KEY] Generating P-384 key pair...")
        self.private_key = ec.generate_private_key(ec.SECP384R1())
        self.public_key = self.private_key.public_key()
        self.shared_key = None
        self.aes = None
    def get_public_key_bytes(self):
        return self.public_key.public_numbers().x.to_bytes(48, 'big') + \
               self.public_key.public_numbers().y.to_bytes(48, 'big')
    def compute_shared_key(self, peer_public_key_bytes):
        try:
            x = int.from_bytes(peer_public_key_bytes[:48], 'big')
            y = int.from_bytes(peer_public_key_bytes[48:], 'big')
            peer_public_numbers = ec.EllipticCurvePublicNumbers(x, y, ec.SECP384R1())
            peer_public_key = peer_public_numbers.public_key()
            shared_secret = self.private_key.exchange(ec.ECDH(), peer_public_key)
            derived_key = HKDF(
                algorithm=hashes.SHA256(),
                length=32,
                salt=None,
                info=b'terminal_chat_v3_p384'
            ).derive(shared_secret)
            self.aes = AESGCM(derived_key)
            print("[SUCCESS] Key exchange succeeded! Using P-384 + AES-256-GCM.")
            return True
        except Exception as e:
            print(f"[ERROR] Key exchange failed: {e}")
            return False
    def encrypt_message(self, message):
        if not self.aes:
            return None
        nonce = secrets.token_bytes(12)
        ciphertext = self.aes.encrypt(nonce, message.encode('utf-8'), None)
        return nonce + ciphertext
    def decrypt_message(self, encrypted_data):
        if not self.aes or len(encrypted_data) < 12:
            return None
        nonce = encrypted_data[:12]
        ciphertext = encrypted_data[12:]
        try:
            plaintext = self.aes.decrypt(nonce, ciphertext, None)
            return plaintext.decode('utf-8')
        except:
            return None


class TerminalChat:
    def __init__(self):
        self.crypto = SimpleCrypto()
        self.socket = None
        self.connected = False
        self.server_socket = None
        self.running = True
        self.receive_thread = None
        self.server_thread = None
        self.print_header()
    def print_header(self):
        print("=" * 60)
        print("  Minimal Encrypted Chat v3.0 - Stable Version")
        print("  P-384 ECDH + AES-256-GCM Encryption")
        print("  IPv6 P2P Direct Connection (All Scenarios Supported)")
        print("=" * 60)
    def get_my_ipv6(self):
        try:
            s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
            s.connect(("2001:4860:4860::8888", 80))
            my_ipv6 = s.getsockname()[0]
            s.close()
            return my_ipv6
        except:
            return "::1"
    def start_server(self):
        def server_thread():
            try:
                self.server_socket = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
                self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                self.server_socket.bind(('::', 8888))
                self.server_socket.listen(1)
                server_addr = self.server_socket.getsockname()
                print(f"[LISTEN] Listening on port {server_addr[1]} (all IPv6 interfaces), waiting for connection...")
                while self.running and not self.connected:
                    try:
                        self.server_socket.settimeout(1.0)
                        client_socket, addr = self.server_socket.accept()
                        if not self.connected:
                            print(f"[CONNECT] Connection from {addr[0]}:{addr[1]}")
                            self.socket = client_socket
                            threading.Thread(target=self.handle_connection, args=(True,), daemon=True).start()
                            return
                        else:
                            client_socket.close()
                    except socket.timeout:
                        continue
                    except Exception as e:
                        if self.running:
                            print(f"[ERROR] Accept failed: {e}")
                        break
            except Exception as e:
                print(f"[ERROR] Server start failed: {e}")
        self.server_thread = threading.Thread(target=server_thread, daemon=True)
        self.server_thread.start()
    def connect_to_peer(self, address):
        try:
            print(f"[CONNECT] Connecting to {address}...")
            if address in ["::1", "127.0.0.1"] or address.startswith("240e:"):
                time.sleep(0.5)
            self.socket = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
            self.socket.connect((address, 8888))
            local_addr = self.socket.getsockname()
            print(f"[LOCAL] Using port {local_addr[1]} to connect to {address}:8888")
            threading.Thread(target=self.handle_connection, args=(False,), daemon=True).start()
            return True
        except Exception as e:
            print(f"[ERROR] Connect failed: {e}")
            return False
    def handle_connection(self, is_server):
        try:
            print("[HANDSHAKE] Starting key exchange...")
            my_public_key = self.crypto.get_public_key_bytes()
            self.socket.settimeout(5.0)
            if is_server:
                peer_public_key = self.socket.recv(96)
                self.socket.send(my_public_key)
            else:
                self.socket.send(my_public_key)
                peer_public_key = self.socket.recv(96)
            if self.crypto.compute_shared_key(peer_public_key):
                self.connected = True
                local_addr = self.socket.getsockname()
                peer_addr = self.socket.getpeername()
                print(f"[CONNECT] Local: {local_addr[0]}:{local_addr[1]} <-> Peer: {peer_addr[0]}:{peer_addr[1]}")
                print("[SUCCESS] Connection established! You can chat securely now.")
                print("[INFO] Type your message and press Enter to send, '/quit' to exit.")
                print("-" * 60)
                self.socket.settimeout(None)
                self.start_receive_thread()
            else:
                print("[ERROR] Key exchange failed")
        except Exception as e:
            print(f"[ERROR] Connection handling failed: {e}")
            self.connected = False
    def start_receive_thread(self):
        def receive_thread():
            while self.connected and self.running:
                try:
                    self.socket.settimeout(None)
                    length_data = self.socket.recv(4)
                    if not length_data:
                        print("[INFO] Peer closed the connection.")
                        break
                    msg_length = int.from_bytes(length_data, 'big')
                    if msg_length <= 0 or msg_length > 1000000:
                        print(f"[WARN] Abnormal message length: {msg_length}")
                        break
                    encrypted_data = b''
                    while len(encrypted_data) < msg_length:
                        chunk = self.socket.recv(msg_length - len(encrypted_data))
                        if not chunk:
                            print("[INFO] Incomplete message received.")
                            break
                        encrypted_data += chunk
                    if len(encrypted_data) != msg_length:
                        print(f"[WARN] Message length mismatch: expected {msg_length}, got {len(encrypted_data)}")
                        continue
                    message = self.crypto.decrypt_message(encrypted_data)
                    if message:
                        timestamp = time.strftime("%H:%M:%S")
                        print(f"\r[{timestamp}] Peer: {message}")
                        print("Me: ", end="", flush=True)
                    else:
                        print("[WARN] Message decryption failed")
                except Exception as e:
                    print(f"[ERROR] Receiving message failed: {e}")
                    break
            if self.connected:
                self.connected = False
                print(f"\n[DISCONNECT] Connection closed.")
        self.receive_thread = threading.Thread(target=receive_thread, daemon=True)
        self.receive_thread.start()
    def send_message(self, message):
        if not self.connected:
            return False
        try:
            encrypted_data = self.crypto.encrypt_message(message)
            if encrypted_data:
                self.socket.send(len(encrypted_data).to_bytes(4, 'big'))
                self.socket.send(encrypted_data)
                return True
            else:
                print("[ERROR] Message encryption failed")
                return False
        except Exception as e:
            print(f"[ERROR] Sending failed: {e}")
            return False
    def input_loop(self):
        while self.connected and self.running:
            try:
                message = input("Me: ").strip()
                if message == '/quit':
                    print("[EXIT] Exiting...")
                    break
                elif message:
                    if self.send_message(message):
                        timestamp = time.strftime("%H:%M:%S")
                        print(f"\r[{timestamp}] Me: {message}")
                        print("Me: ", end="", flush=True)
                    else:
                        print("[ERROR] Sending failed, connection may be closed.")
                        break
            except (KeyboardInterrupt, EOFError):
                print("\n[INTERRUPT] User interrupted.")
                break
            except Exception as e:
                print(f"[ERROR] Input loop failed: {e}")
                break
        self.cleanup()
    def cleanup(self):
        print("[CLEANUP] Cleaning up resources...")
        self.running = False
        self.connected = False
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
        if self.server_socket:
            try:
                self.server_socket.close()
            except:
                pass
        print("[DONE] Program exited, chat history cleared.")
    def run(self):
        try:
            my_ipv6 = self.get_my_ipv6()
            print(f"[ADDRESS] My IPv6 address: {my_ipv6}")
            print("[INFO] The peer needs to use this address to connect to you.")
            self.start_server()
            time.sleep(0.5)
            print("\nPlease choose:")
            print("1. Wait for peer to connect to me")
            print("2. Connect to peer")
            while True:
                choice = input("Enter choice (1/2): ").strip()
                if choice == '1':
                    print("[WAIT] Waiting for peer to connect...")
                    while not self.connected and self.running:
                        time.sleep(0.1)
                    if self.connected:
                        self.input_loop()
                    break
                elif choice == '2':
                    peer_address = input("Enter peer's IPv6 address: ").strip()
                    if peer_address:
                        if self.connect_to_peer(peer_address):
                            timeout = 10
                            while not self.connected and self.running and timeout > 0:
                                time.sleep(0.1)
                                timeout -= 1
                            if self.connected:
                                self.input_loop()
                                break
                            else:
                                print("[ERROR] Connection timeout, please try again.")
                        else:
                            print("[ERROR] Connection failed, please try again.")
                else:
                    print("Please enter 1 or 2.")
        except KeyboardInterrupt:
            print("\n\n[EXIT] Program interrupted.")
            self.cleanup()
def main():
    try:
        chat = TerminalChat()
        chat.run()
    except Exception as e:
        print(f"Program error: {e}")

if __name__ == "__main__":
    main() 